from src.agents.outline_writer import outlineWriter
from src.database import PaperDatabase
from src.agents.writer import subsectionWriter
import os

'''
# 创建 Local 数据库实例
db = PaperDatabase(
    db_type="local",
    local_config={
        "db_path": "./database",
        "embedding_model": "./embedding_models/AI-ModelScope/nomic-embed-text-v1",
        "device": "cuda"  # 或 "cpu"
    }
)
'''
# 创建 Atomgit 数据库实例
db = PaperDatabase(
    db_type="atomgit",
    atomgit_config={
        "base_url": "http://180.184.65.98:38880/atomgit"
    }
)

model_name = 'glm-4-air'
api_key = 'XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX'
api_url = "https://open.bigmodel.cn/api/paas/v4"

topic = "Research Progress and Development Trends in Knowledge Graphs"
saving_path= './output/'




outline_writer = outlineWriter(model=model_name, api_key=api_key, api_url=api_url, database=db)
subsection_writer = subsectionWriter(model=model_name, api_key=api_key, api_url=api_url, database=db)

print("--- 准备写outline ---")
print('--- 论文检索中 ---')
outline = outline_writer.draft_outline(topic)
print("outline",outline)
    
raw_survey, raw_survey_with_references, raw_references, refined_survey, refined_survey_with_references, refined_references = subsection_writer.write(topic=topic, outline=outline)

if not os.path.exists(saving_path):
    os.mkdir(saving_path)

with open(f'{saving_path}/{topic}.md', 'a+') as f:
    f.write(refined_survey_with_references)